use crate::mimalloc::allocator::MiMalloc;
use crate::mimalloc::utils::{align_up, ChunkCategory, MemoryBlock};
use crate::mimalloc::PAGE_SIZE;
use std::alloc::{GlobalAlloc, Layout};
use std::ptr::{null_mut, NonNull};
use std::sync::atomic::{AtomicPtr, Ordering};

impl MiMalloc {
    /// 创建新的分配器实例
    pub const fn new() -> Self {
        Self {
            free_lists: [
                AtomicPtr::new(null_mut()),
                AtomicPtr::new(null_mut()),
                AtomicPtr::new(null_mut()),
            ],
        }
    }

    /// 根据大小确定内存块类别
    pub(crate) fn categorize_chunk(size: usize) -> ChunkCategory {
        match size {
            0..=64 => ChunkCategory::Small,
            65..=256 => ChunkCategory::Medium,
            _ => ChunkCategory::Large,
        }
    }

    /// 获取对应大小类别的空闲列表
    pub(crate) fn get_free_list(&self, category: ChunkCategory) -> &AtomicPtr<MemoryBlock> {
        &self.free_lists[category as usize]
    }

    /// 从空闲列表分配内存
    pub(crate) unsafe fn allocate_from_free_list(
        &self,
        free_list: &AtomicPtr<MemoryBlock>,
        layout: Layout,
    ) -> *mut u8 {
        let aligned_size = align_up(layout.size(), layout.align());

        let mut current = free_list.load(Ordering::SeqCst);

        while !current.is_null() {
            let block = &mut *current;
            if block.size >= aligned_size {
                // 尝试原子地更新列表头
                let next = block.next.map_or(null_mut(), |n| n.as_ptr());
                if free_list
                    .compare_exchange(current, next, Ordering::SeqCst, Ordering::SeqCst)
                    .is_ok()
                {
                    block.is_free = false;

                    // 如果剩余空间足够大，可以进行分割
                    if block.size > aligned_size + std::mem::size_of::<MemoryBlock>() {
                        let remainder = block.size - aligned_size;
                        let remainder_block =
                            (current as *mut u8).add(aligned_size) as *mut MemoryBlock;

                        // 初始化剩余块的元数据
                        (*remainder_block).size = remainder;
                        (*remainder_block).is_free = true;
                        (*remainder_block).next = block.next;

                        // 将剩余块放回空闲列表
                        self.free_to_list(
                            remainder_block as *mut u8,
                            Layout::from_size_align(remainder, layout.align()).unwrap(),
                        );
                    }

                    return current as *mut u8;
                }
                // 如果交换失败，重新开始
                current = free_list.load(Ordering::SeqCst);
            } else {
                current = block.next.map_or(null_mut(), |n| n.as_ptr());
            }
        }

        null_mut()
    }

    /// 系统内存分配
    pub(crate) unsafe fn system_alloc(size: usize) -> *mut u8 {
        #[cfg(target_os = "linux")]
        {
            // Linux 平台使用 mmap
            let ptr = libc::mmap(
                std::ptr::null_mut(),
                size,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
                -1,
                0,
            );

            if ptr == libc::MAP_FAILED {
                return std::ptr::null_mut();
            }

            ptr as *mut u8
        }

        #[cfg(not(target_os = "linux"))]
        {
            // 非 Linux 平台使用标准分配器
            std::alloc::System.alloc(Layout::from_size_align(size, PAGE_SIZE).unwrap())
        }
    }

    /// 将内存块释放到空闲列表
    pub(crate) unsafe fn free_to_list(&self, ptr: *mut u8, layout: Layout) {
        let category = Self::categorize_chunk(layout.size());
        let free_list = self.get_free_list(category);

        let block = ptr as *mut MemoryBlock;
        (*block).is_free = true;
        (*block).size = layout.size();

        // 原子地更新列表头
        loop {
            let current = free_list.load(Ordering::SeqCst);

            // 使用 Option<NonNull> 安全地处理链接
            (*block).next = NonNull::new(current);

            // 尝试将新块放在列表头
            match free_list.compare_exchange_weak(
                current,
                block,
                Ordering::SeqCst,
                Ordering::SeqCst,
            ) {
                Ok(_) => break,
                Err(_) => continue,
            }
        }
    }
}

unsafe impl Sync for MiMalloc {}
unsafe impl Send for MiMalloc {}
